#include "direct_solver.hpp"

template<typename idx_t, typename data_t, typename res_t>
CSR_sparseMat<idx_t, data_t, res_t>::CSR_sparseMat(int root_pid, idx_t m, idx_t * Ai, idx_t * Aj, data_t * Av):
    root_pid(root_pid), nrows(m), row_ptr(Ai), col_idx(Aj), vals(Av)
{
    // 检查升序
#ifdef DEBUG
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    for (idx_t i = 0; i < m; i++) {
        idx_t pbeg = Ai[i], pend = Ai[i+1];
        for (idx_t p = pbeg + 1; p < pend; p++) {
            if (Aj[p-1] > Aj[p]) {
                printf("Proc %d: Aj not increasing! i %d j %d %d\n", my_pid, i, Aj[p-1], Aj[p]);
                MPI_Abort(MPI_COMM_WORLD, -5061);
            }
        }
    }
#endif
}

template<typename idx_t, typename data_t, typename res_t>
CSR_sparseMat<idx_t, data_t, res_t>::~CSR_sparseMat()
{
    if (row_ptr) { delete [] row_ptr; row_ptr = nullptr; }
    if (col_idx) { delete [] col_idx; col_idx = nullptr; }
    if (vals   ) { delete [] vals;    vals    = nullptr; }
#ifdef __x86_64__
    phase = -1;// Release internal memory.
    PARDISO (pt, &maxfct, &mnum, &mtype, &phase,
            &nrows, &ddum, row_ptr, col_idx, &idum, &nrhs,
            iparm, &msglvl, &ddum, &ddum, &error);
#endif
}

template<typename idx_t, typename data_t, typename res_t>
bool CSR_sparseMat<idx_t, data_t, res_t>::initSolver(const int nthreads, bool stable, int isComplex)
{
#ifdef __aarch64__
    // Generally, only need to specify the required fields, such as VALUE_TYPE and INDEX_TYPE
    TydssSolverOptions opt;
    opt.fieldMask = TYDSS_SOLVER_OPTIONS_VALUE_TYPE |
                    TYDSS_SOLVER_OPTIONS_INDEX_TYPE |
                    TYDSS_SOLVER_OPTIONS_NTHREADS |
                    TYDSS_SOLVER_OPTIONS_STABLE;
    if (isComplex == 0) {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64 : TYDSS_VALUE_FP32;
    } else {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64C : TYDSS_VALUE_FP32C;
    }
    opt.indexType = sizeof(idx_t) == 4 ? TYDSS_INDEX_INT32: TYDSS_INDEX_INT64;
    opt.numThreads = nthreads;
    opt.stable = stable;
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    if (my_pid == root_pid) fprintf(stdout, "Enable stable: %s\n", stable == 0 ? "false" : "true");
    if (solver.Init(&opt) != TYDSS_OK) {
        return false;
    }
#elif defined(__x86_64__)
    for (MKL_INT i = 0; i < 64; i++) iparm[i] = 0;
    iparm[0] = 1;// no solver default
    iparm[1] = 2;// Fill-in reordering from METIS
    iparm[3] = 0;// No iterative-direct algorithm
    iparm[4] = 0;// No user fill-in reducing permutation
    iparm[5] = 0;// Write solution into x
    iparm[7] = 0;// Max numbers of iterative refinement steps
    iparm[9] = 13;// Perturb the pivot elements with 1E-13
    iparm[10] = 1;// Use nonsymmetric permutation and scaling MPS
    iparm[12] = 0;// Maximum weighted matching algorithm is switched-off (default for symmetric). Try iparm[12] = 1 in case of inappropriate accuracy
    iparm[13] = 0;// Output: Number of perturbed pivots
    iparm[17] = -1;// Output: Number of nonzeros in the factor LU
    iparm[18] = -1;// Output: Mflops for LU factorization
    iparm[19] = 0;// Output: Numbers of CG Iterations
    iparm[27] = sizeof(data_t) == 4 ? 1 : 0;// 0 for double, 1 for float
    iparm[34] = 1;// PARDISO use C-style indexing for ia and ja arrays */
    maxfct = 20;// Maximum number of numerical factorizations. */
    mnum = 1;// Which factorization to use.
    msglvl = 0;// Print statistical information in file */
    error = 0;// Initialize error flag
    for (MKL_INT i = 0; i < 64; i++) pt[i] = 0;
#endif
    return true;
}

template<typename idx_t, typename data_t, typename res_t>
bool CSR_sparseMat<idx_t, data_t, res_t>::initA(int isComplex, PATTERN pattern)
{
    int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
#ifdef __aarch64__
    TydssMatrixStore store;
    store.nrows = nrows;
    store.ncols = nrows;
    store.type = TYDSS_MATRIX_STORE_CSR;
    store.csr.offsets = row_ptr;
    store.csr.indices = col_idx;
    store.csr.values = vals;

    TydssMatrixOptions opt;
    opt.fieldMask = TYDSS_MATRIX_OPTIONS_VALUE_TYPE |
                    TYDSS_MATRIX_OPTIONS_INDEX_TYPE |
                    TYDSS_MATRIX_OPTIONS_PATTERN |
                    TYDSS_MATRIX_OPTIONS_ATTR;
    if (isComplex == 0) {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64 : TYDSS_VALUE_FP32;
    } else {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64C : TYDSS_VALUE_FP32C;
    }
    opt.indexType = sizeof(idx_t) == 4 ? TYDSS_INDEX_INT32: TYDSS_INDEX_INT64;
    
    switch (pattern) {
        case SYMMETRIC_POSITIVE_DEFINITE: // symmetric positive definite
            opt.pattern = TYDSS_MATRIX_PATTERN_SYM;
            opt.attr = TYDSS_MATRIX_ATTR_DEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is symmetric positive definite.\n");
            break;
        case SYMMETRIC: // symmetric indefinite
            opt.pattern = TYDSS_MATRIX_PATTERN_SYM;
            opt.attr = TYDSS_MATRIX_ATTR_INDEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is symmetric indefinite.\n");
            break;
        case STRUCTURAL_SYMMETRIC: // structural symmetric
            opt.pattern = TYDSS_MATRIX_PATTERN_STRUCT_SYM;
            opt.attr = TYDSS_MATRIX_ATTR_INDEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is structural symmetric.\n");
            break;
        case NONSYMMETRIC: // non symmetric
            opt.pattern = TYDSS_MATRIX_PATTERN_NON_SYM;
            opt.attr = TYDSS_MATRIX_ATTR_INDEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is non symmetric.\n");
            break;
        case C_HERMITIAN_POSITIVE_DEFINITE: // hermitian positive definite
            opt.pattern = TYDSS_MATRIX_PATTERN_HERMITIAN;
            opt.attr = TYDSS_MATRIX_ATTR_DEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is hermitian positive definite.\n");
            break;
        case C_HERMITIAN: // hermitian indefinite
            opt.pattern = TYDSS_MATRIX_PATTERN_HERMITIAN;
            opt.attr = TYDSS_MATRIX_ATTR_INDEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is hermitian indefinite.\n");
            break;
        default:
            opt.pattern = TYDSS_MATRIX_PATTERN_NON_SYM;
            opt.attr = TYDSS_MATRIX_ATTR_INDEFINITE;
            if (my_pid == root_pid) printf("Tydss: Matrix is non symmetric.\n");
            break;
    }

    if (A.Init(store, &opt) != TYDSS_OK) {
        printf("Failed to init A\n");
        return false;
    }
#elif defined(__x86_64__)
    assert(isComplex == 0);// 还没支持
    switch (pattern) {
        case SYMMETRIC_POSITIVE_DEFINITE: // symmetric positive definite
            mtype = 2;
            if (my_pid == root_pid) printf("Pardiso: Matrix is real symmetric positive definite.\n");
            break;
        case SYMMETRIC: // symmetric indefinite
            mtype = -2;
            if (my_pid == root_pid) printf("Pardiso: Matrix is real symmetric indefinite.\n");
            break;
        case STRUCTURAL_SYMMETRIC: // structural symmetric
            mtype = 1;
            if (my_pid == root_pid) printf("Pardiso: Matrix is real structural symmetric.\n");
            break;
        case NONSYMMETRIC: // non symmetric
            mtype = 11;
            if (my_pid == root_pid) printf("Pardiso: Matrix is real non symmetric.\n");
            break;
        case C_HERMITIAN_POSITIVE_DEFINITE: // hermitian positive definite
            mtype = 4;
            if (my_pid == root_pid) printf("Pardiso: Matrix is complex hermitian positive definite.\n");
            break;
        case C_HERMITIAN: // hermitian indefinite
            mtype = -4;
            if (my_pid == root_pid) printf("Pardiso: Matrix is complex hermitian indefinite.\n");
            break;
        case C_SYMMETRIC:
            mtype = 6;
            if (my_pid == root_pid) printf("Pardiso: Matrix is complex symmetric.\n");
            break;
        case C_STRUCTURAL_SYMMETRIC:
            mtype = 3;
            if (my_pid == root_pid) printf("Pardiso: Matrix is complex structural symmetric.\n");
            break;
        case C_NONSYMMETRIC:
            mtype = 13;
            if (my_pid == root_pid) printf("Pardiso: Matrix is complex non symmetric.\n");
            break;
        default:
            mtype = 11;
            if (my_pid == root_pid) printf("Pardiso: Matrix is real non symmetric.\n");
            break;
    }
#endif
    return true;
}

template<typename idx_t, typename data_t, typename res_t>
bool CSR_sparseMat<idx_t, data_t, res_t>::initB(const idx_t nrhs, int isComplex)
{
#ifdef __aarch64__
    TydssMatrixStore store;
    store.nrows = nrows;
    store.ncols = nrhs;
    store.type = TYDSS_MATRIX_STORE_DENSE_COL_MAJOR;
    store.dense.values = nullptr;

    TydssMatrixOptions opt;
    opt.fieldMask = TYDSS_MATRIX_OPTIONS_VALUE_TYPE |
                    TYDSS_MATRIX_OPTIONS_INDEX_TYPE;
    if (isComplex == 0) {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64 : TYDSS_VALUE_FP32;
    } else {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64C : TYDSS_VALUE_FP32C;
    }
    opt.indexType = sizeof(idx_t) == 4 ? TYDSS_INDEX_INT32: TYDSS_INDEX_INT64;
    if (B.Init(store, &opt) != TYDSS_OK) {
        printf("Failed to init B\n");
        return false;
    }
#elif defined(__x86_64__)
    assert(nrhs == 1);
#endif
    return true;
}

template<typename idx_t, typename data_t, typename res_t>
bool CSR_sparseMat<idx_t, data_t, res_t>::initX(const idx_t nrhs, int isComplex)
{
#ifdef __aarch64__
    /* Matrix representing multi solutions */
    TydssMatrixStore store;
    store.nrows = nrows;
    store.ncols = nrhs;
    store.type = TYDSS_MATRIX_STORE_DENSE_COL_MAJOR;
    store.dense.values = nullptr;

    TydssMatrixOptions opt;
    opt.fieldMask = TYDSS_MATRIX_OPTIONS_VALUE_TYPE |
                    TYDSS_MATRIX_OPTIONS_INDEX_TYPE;
    if (isComplex == 0) {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64 : TYDSS_VALUE_FP32;
    } else {
        opt.valueType = sizeof(data_t) == 8 ? TYDSS_VALUE_FP64C : TYDSS_VALUE_FP32C;
    }
    opt.indexType = sizeof(idx_t) == 4 ? TYDSS_INDEX_INT32: TYDSS_INDEX_INT64;
    if (X.Init(store, &opt) != TYDSS_OK) {
        printf("Failed to init X\n");
        return false;
    }
#elif defined(__x86_64__)
    assert(nrhs == 1);
#endif
    return true;
}

template<typename idx_t, typename data_t, typename res_t>
void CSR_sparseMat<idx_t, data_t, res_t>::analyse(const idx_t nrhs, int swm, int wm,
        const int nthreads, const int nthreadsFR, const int dynamic)
{
#ifdef __aarch64__
    TydssAnalyseOptions opt;
    opt.fieldMask = TYDSS_ANALYSE_OPTIONS_WM |
                    TYDSS_ANALYSE_OPTIONS_SWM |
                    TYDSS_ANALYSE_OPTIONS_FR |
                    TYDSS_ANALYSE_OPTIONS_NTHREADS_FR |
                    TYDSS_ANALYSE_OPTIONS_BATCH_NRHS |
            TYDSS_ANALYSE_OPTIONS_PIVOT;
    opt.swm = swm == 1 ? TYDSS_SWM_MC64 : TYDSS_SWM_NONE;
    opt.wm = wm == 1 ? TYDSS_WM_MC64 : TYDSS_WM_NONE;
    opt.fr = TYDSS_FR_METIS;
    opt.numThreadsFR = nthreadsFR;
    opt.batchNRHS = nrhs;
    if (dynamic == 1) {
        opt.pivotAlgo = TYDSS_PIVOT_DYNAMIC;
    }
    /* consistent with 0.9.0 */
    idx_t maxAnalyseThreads = 32;
    if (nthreads > maxAnalyseThreads) {
        opt.fieldMask |= TYDSS_ANALYSE_OPTIONS_NTHREADS_ANALYSE;
        opt.numThreadsAnalyse = maxAnalyseThreads;
    }
    idx_t maxApplyThreads = 2 * 32;
    if (nrhs == 1 && nthreads > maxApplyThreads) {
        opt.fieldMask |= TYDSS_ANALYSE_OPTIONS_NTHREADS_APPLY;
        opt.numThreadsApply = maxApplyThreads;
    }
    TydssStatus rtn = solver.Analyse(A, &opt);
    assert(rtn == TYDSS_OK);
#elif defined(__x86_64__)
    phase = 11;
    PARDISO (pt, &maxfct, &mnum, &mtype, &phase,
            &nrows, vals, row_ptr, col_idx, &idum, &nrhs, iparm, &msglvl, &ddum, &ddum, &error);
    if (error != 0 ) {
        printf ("\nPardiso: ERROR during symbolic factorization: %d", error);
        MPI_Abort(MPI_COMM_WORLD, -20230113);
    }
#endif
    return;
}

template<typename idx_t, typename data_t, typename res_t>
void CSR_sparseMat<idx_t, data_t, res_t>::factorize(double pert)
{
#ifdef __aarch64__
    TydssFactorizeOptions fopt;
    fopt.fieldMask = TYDSS_FACTORIZE_OPTIONS_PERT;
    fopt.pert = pert;
    TydssStatus rtn = solver.Factorize(A, &fopt);
    assert(rtn == TYDSS_OK);
#elif defined(__x86_64__)
    iparm[13] = pert;
    phase = 22;
    PARDISO (pt, &maxfct, &mnum, &mtype, &phase,
        &nrows, vals, row_ptr, col_idx, &idum, &nrhs, iparm, &msglvl, &ddum, &ddum, &error);
    if (error != 0 ) {
        printf ("\nPardiso: ERROR during numerical factorization: %d", error);
        MPI_Abort(MPI_COMM_WORLD, -20230113);
    }
#endif
    return;
}

template<typename idx_t, typename data_t, typename res_t>
void CSR_sparseMat<idx_t, data_t, res_t>::apply(res_t* b, res_t* x, int refine)
{
#ifdef __aarch64__
    B.SetValues(b);
    X.SetValues(x);
    TydssApplyOptions options;
    options.fieldMask = TYDSS_APPLY_OPTIONS_REFINE_METHOD;
    options.refineMethod = refine == 1 ? TYDSS_REFINE_NAIVE : TYDSS_REFINE_NONE;
    TydssStatus rtn = solver.Apply(B, X, &options);
    assert(rtn == TYDSS_OK);
#elif defined(__x86_64__)
    phase = 33;
    PARDISO (pt, &maxfct, &mnum, &mtype, &phase,
            &nrows, vals, row_ptr, col_idx, &idum, &nrhs, iparm, &msglvl, b, x, &error);
    if (error != 0 ) {
        printf ("\nPardiso: ERROR during solution: %d", error);
        MPI_Abort(MPI_COMM_WORLD, -20230113);
    }
#endif
    return;
}

template<typename idx_t, typename data_t, typename res_t>
void CSR_sparseMat<idx_t, data_t, res_t>::fprint_COO(const char * filename) const
{
    FILE * fp = fopen(filename, "w+");
    for (idx_t i = 0; i < nrows; i++) {
        idx_t pbeg = row_ptr[i], pend = row_ptr[i+1];
        for (idx_t p = pbeg; p < pend; p++)
            fprintf(fp, "%d %d %.5e\n", i, col_idx[p], vals[p]);
    }
    fclose(fp);
}

template class CSR_sparseMat<int, double, double>;
template class CSR_sparseMat<int, float , float >;
